#!/usr/bin/env python3
"""
SQLite到MySQL数据迁移脚本
将现有的SQLite数据库迁移到MySQL数据库
"""

import os
import sys
import sqlite3
import pymysql
import logging
from datetime import datetime
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from config.mysql_config import MySQLConfig

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class DatabaseMigrator:
    """数据库迁移器"""
    
    def __init__(self):
        self.mysql_config = MySQLConfig()
        self.mysql_config.validate_config()
        
        # SQLite数据库路径
        self.sqlite_path = project_root / "instance" / "dingdingrobot.db"
        
        # 表映射（SQLite表名 -> MySQL表名）
        self.table_mapping = {
            'user': 'user',
            'ding_talk_robot': 'ding_talk_robot',
            'scheduled_task': 'scheduled_task',
            'task_execution_log': 'task_execution_log',
            'custom_script': 'custom_script',
            'custom_task': 'custom_task',
            'background_installation_task': 'background_installation_task',
            'package_installation_log': 'package_installation_log',
            'script': 'script',
            'task_robot_association': 'task_robot_association'
        }
    
    def connect_sqlite(self):
        """连接SQLite数据库"""
        if not self.sqlite_path.exists():
            raise FileNotFoundError(f"SQLite数据库文件不存在: {self.sqlite_path}")
        
        conn = sqlite3.connect(str(self.sqlite_path))
        conn.row_factory = sqlite3.Row  # 使结果可以按列名访问
        return conn
    
    def connect_mysql(self):
        """连接MySQL数据库"""
        return pymysql.connect(
            host=self.mysql_config.MYSQL_HOST,
            port=self.mysql_config.MYSQL_PORT,
            user=self.mysql_config.MYSQL_USER,
            password=self.mysql_config.MYSQL_PASSWORD,
            database=self.mysql_config.MYSQL_DATABASE,
            charset=self.mysql_config.MYSQL_CHARSET,
            autocommit=False
        )
    
    def get_sqlite_tables(self):
        """获取SQLite数据库中的所有表"""
        sqlite_conn = self.connect_sqlite()
        cursor = sqlite_conn.cursor()
        
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'")
        tables = [row[0] for row in cursor.fetchall()]
        
        sqlite_conn.close()
        return tables
    
    def get_table_schema(self, table_name):
        """获取表结构"""
        sqlite_conn = self.connect_sqlite()
        cursor = sqlite_conn.cursor()
        
        cursor.execute(f"PRAGMA table_info({table_name})")
        columns = cursor.fetchall()
        
        sqlite_conn.close()
        return columns
    
    def get_table_data(self, table_name):
        """获取表数据"""
        sqlite_conn = self.connect_sqlite()
        cursor = sqlite_conn.cursor()
        
        cursor.execute(f"SELECT * FROM {table_name}")
        data = cursor.fetchall()
        
        sqlite_conn.close()
        return data
    
    def check_mysql_table_exists(self, table_name):
        """检查MySQL表是否存在"""
        mysql_conn = self.connect_mysql()
        cursor = mysql_conn.cursor()
        
        cursor.execute(f"""
            SELECT COUNT(*) 
            FROM information_schema.tables 
            WHERE table_schema = '{self.mysql_config.MYSQL_DATABASE}' 
            AND table_name = '{table_name}'
        """)
        
        exists = cursor.fetchone()[0] > 0
        mysql_conn.close()
        return exists
    
    def migrate_table_data(self, sqlite_table, mysql_table):
        """迁移单个表的数据"""
        logger.info(f"开始迁移表: {sqlite_table} -> {mysql_table}")
        
        # 检查MySQL表是否存在
        if not self.check_mysql_table_exists(mysql_table):
            logger.warning(f"MySQL表 {mysql_table} 不存在，跳过迁移")
            return False
        
        # 获取SQLite表数据
        data = self.get_table_data(sqlite_table)
        if not data:
            logger.info(f"表 {sqlite_table} 没有数据，跳过迁移")
            return True
        
        # 获取表结构
        columns = self.get_table_schema(sqlite_table)
        column_names = [col[1] for col in columns]
        
        # 连接MySQL
        mysql_conn = self.connect_mysql()
        cursor = mysql_conn.cursor()
        
        try:
            # 清空目标表（可选）
            cursor.execute(f"DELETE FROM {mysql_table}")
            logger.info(f"清空MySQL表 {mysql_table}")
            
            # 构建插入语句
            placeholders = ', '.join(['%s'] * len(column_names))
            insert_sql = f"INSERT INTO {mysql_table} ({', '.join(column_names)}) VALUES ({placeholders})"
            
            # 批量插入数据
            batch_size = 1000
            total_rows = len(data)
            
            for i in range(0, total_rows, batch_size):
                batch = data[i:i + batch_size]
                batch_data = [tuple(row) for row in batch]
                
                cursor.executemany(insert_sql, batch_data)
                mysql_conn.commit()
                
                logger.info(f"已迁移 {min(i + batch_size, total_rows)}/{total_rows} 行数据")
            
            logger.info(f"表 {sqlite_table} 迁移完成，共迁移 {total_rows} 行数据")
            return True
            
        except Exception as e:
            mysql_conn.rollback()
            logger.error(f"迁移表 {sqlite_table} 失败: {e}")
            return False
        finally:
            mysql_conn.close()
    
    def verify_migration(self):
        """验证迁移结果"""
        logger.info("开始验证迁移结果...")
        
        sqlite_conn = self.connect_sqlite()
        mysql_conn = self.connect_mysql()
        
        sqlite_cursor = sqlite_conn.cursor()
        mysql_cursor = mysql_conn.cursor()
        
        verification_results = {}
        
        try:
            for sqlite_table, mysql_table in self.table_mapping.items():
                # 检查表是否存在
                sqlite_cursor.execute(f"SELECT name FROM sqlite_master WHERE type='table' AND name='{sqlite_table}'")
                sqlite_exists = sqlite_cursor.fetchone() is not None
                
                mysql_cursor.execute(f"""
                    SELECT COUNT(*) 
                    FROM information_schema.tables 
                    WHERE table_schema = '{self.mysql_config.MYSQL_DATABASE}' 
                    AND table_name = '{mysql_table}'
                """)
                mysql_exists = mysql_cursor.fetchone()[0] > 0
                
                if not sqlite_exists:
                    verification_results[sqlite_table] = "SQLite表不存在"
                    continue
                
                if not mysql_exists:
                    verification_results[sqlite_table] = "MySQL表不存在"
                    continue
                
                # 比较行数
                sqlite_cursor.execute(f"SELECT COUNT(*) FROM {sqlite_table}")
                sqlite_count = sqlite_cursor.fetchone()[0]
                
                mysql_cursor.execute(f"SELECT COUNT(*) FROM {mysql_table}")
                mysql_count = mysql_cursor.fetchone()[0]
                
                if sqlite_count == mysql_count:
                    verification_results[sqlite_table] = f"✅ 成功 ({sqlite_count} 行)"
                else:
                    verification_results[sqlite_table] = f"❌ 失败 (SQLite: {sqlite_count}, MySQL: {mysql_count})"
            
            # 输出验证结果
            logger.info("迁移验证结果:")
            for table, result in verification_results.items():
                logger.info(f"  {table}: {result}")
            
            return verification_results
            
        finally:
            sqlite_conn.close()
            mysql_conn.close()
    
    def run_migration(self):
        """执行完整的迁移过程"""
        logger.info("开始数据库迁移...")
        
        # 测试连接
        logger.info("测试数据库连接...")
        try:
            sqlite_conn = self.connect_sqlite()
            sqlite_conn.close()
            logger.info("SQLite连接成功")
        except Exception as e:
            logger.error(f"SQLite连接失败: {e}")
            return False
        
        try:
            mysql_conn = self.connect_mysql()
            mysql_conn.close()
            logger.info("MySQL连接成功")
        except Exception as e:
            logger.error(f"MySQL连接失败: {e}")
            return False
        
        # 获取SQLite表列表
        sqlite_tables = self.get_sqlite_tables()
        logger.info(f"发现SQLite表: {sqlite_tables}")
        
        # 迁移数据
        success_count = 0
        total_count = 0
        
        for sqlite_table in sqlite_tables:
            if sqlite_table in self.table_mapping:
                mysql_table = self.table_mapping[sqlite_table]
                total_count += 1
                
                if self.migrate_table_data(sqlite_table, mysql_table):
                    success_count += 1
            else:
                logger.warning(f"未知表 {sqlite_table}，跳过迁移")
        
        logger.info(f"迁移完成: {success_count}/{total_count} 个表迁移成功")
        
        # 验证迁移结果
        verification_results = self.verify_migration()
        
        # 统计验证结果
        success_verifications = sum(1 for result in verification_results.values() if "✅" in result)
        total_verifications = len(verification_results)
        
        logger.info(f"验证完成: {success_verifications}/{total_verifications} 个表验证成功")
        
        return success_count == total_count and success_verifications == total_verifications

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='SQLite到MySQL数据迁移工具')
    parser.add_argument('--verify-only', action='store_true', help='仅验证迁移结果，不执行迁移')
    parser.add_argument('--force', action='store_true', help='强制执行迁移（清空目标表）')
    
    args = parser.parse_args()
    
    migrator = DatabaseMigrator()
    
    if args.verify_only:
        logger.info("仅执行验证...")
        verification_results = migrator.verify_migration()
        success_count = sum(1 for result in verification_results.values() if "✅" in result)
        total_count = len(verification_results)
        
        if success_count == total_count:
            logger.info("所有表验证成功！")
            sys.exit(0)
        else:
            logger.error(f"验证失败: {success_count}/{total_count} 个表验证成功")
            sys.exit(1)
    else:
        if not args.force:
            response = input("确定要执行数据迁移吗？这将清空MySQL目标表的现有数据。(y/N): ")
            if response.lower() != 'y':
                logger.info("迁移已取消")
                sys.exit(0)
        
        success = migrator.run_migration()
        sys.exit(0 if success else 1)

if __name__ == '__main__':
    main()